import heapq


class TopKHeap(object):
    def __init__(self, k):
        self.k = k
        self.data = []

    def push(self, elem):
        if len(self.data) < self.k:
            heapq.heappush(self.data, elem)
        else:
            topk_small = self.data[0]
            if elem > topk_small:
                heapq.heapreplace(self.data, elem)

    def top(self):
        return self.data[0]


class Solution:
    def findKthLargest(self, nums, k: int) -> int:
        heap = TopKHeap(k)
        for i in nums:
            heap.push(i)
        return heap.top()


if __name__ == '__main__':
    so = Solution()
    s = [3, 2, 1, 5, 6, 4]
    k = 3
    r = so.findKthLargest(s, k)
    print(r)
